import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOv10
import pymysql
import time
import uuid
import csv
# import mysql.connector

if __name__ == '__main__':

    # 以下参数需要修改成当前训练环境对应的内容，以便进行训练和写入数据库
    cp = "aistudio"
    gpu = "Tesla V100-SXM2-32GB"    
    my_model = '/home/aistudio/work/yolov10-main/ultralytics/cfg/models/v10/yolov10n.yaml'
    my_data = '/home/aistudio/work/yolov10-main/data.yaml'
    imgsz = 640
    epochs = 10
    batch = 64
    workers = 2
    device = '0'
    optimizer = 'Adam'
    close_mosaic = 7
    resume = False
    # uuid_str = str(uuid.uuid4())
    project='runs/train'
    current_time = time.localtime()
    train_datetime = time.strftime("%Y%m%d%H%M%S", current_time)
    name='train_' + str(train_datetime)
    single_cls = False
    cache = True
    sava_dir = project + '/' + name
    start_time = time.time()
    

    # model.load('yolov8n.pt') # 加载预训练权重,改进或者做对比实验时候不建议打开，因为用预训练模型整体精度没有很明显的提升
    model = YOLOv10(model=my_model)
    model.train(data=my_data,
                imgsz=imgsz,
                epochs=epochs,
                batch=batch,
                workers=workers,
                device=device,
                optimizer=optimizer,
                close_mosaic=close_mosaic,
                resume=resume,
                project='runs/train',
                name=name,
                single_cls=single_cls,
                cache=cache,
                )
    # print('running')
    # time.sleep(3)
    # print('stop')

    end_time = time.time()
    # 计算训练结束时间
    time_delta  = end_time - start_time

    hours = int(time_delta // 3600)
    minutes = int((time_delta % 3600) // 60)
    seconds = int(time_delta % 60)
    # 格式化为符合 TIME 类型的字符串
    train_time = f"{hours:02d}:{minutes:02d}:{seconds:02d}"

    # 写入csv文件
    with open('train.csv', 'a', newline='', encoding='utf-8') as csvfile:
        # 创建一个 CSV 写入器对象
        writer = csv.writer(csvfile)
        # 写入单行数据
        writer.writerow([cp, gpu, my_model, my_data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time, train_datetime])

    
    # 建立数据库连接
    # conn = pymysql.connect(
    #     host='chengshaohao.mysql.rds.aliyuncs.com',
    #     user='root',
    #     password='Csh@9560',
    #     database='yolov10',
    #     charset='utf8mb4',
    #     cursorclass=pymysql.cursors.DictCursor
    # )

    # try:
    #     with conn.cursor() as cursor:
    #         # 定义插入数据的 SQL 语句
    #         sql = "INSERT INTO train (cp, gpu, model, data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
    #         # 要插入的数据
    #         val = (cp, gpu, my_model, my_data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time)
    #         # 执行插入语句
    #         cursor.execute(sql, val)
        
    #     # 提交事务，使插入操作生效
    #     conn.commit()
    #     print(cursor.rowcount, "写入完成")
    # except Exception as e:
    #     print(f"Error: {e}")
    # finally:
    #     # 关闭数据库连接
    #     conn.close()



    # try:
    #     conn = mysql.connector.connect(
    #         host='chengshaohao.mysql.rds.aliyuncs.com',
    #         user='root',
    #         password='Csh@9560',
    #         database='yolov10',
    #     )

    #     cursor = conn.cursor()

    #     insert_data_query = "INSERT INTO train (cp, gpu, model, data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
    #     train_data = (cp, gpu, my_model, my_data, imgsz, epochs, batch, workers, device, optimizer, close_mosaic, resume, single_cls, cache, sava_dir, train_time)

    #     cursor.executemany(insert_data_query, user_data)

    #     # 提交更改
    #     conn.commit()
    # except Exception as e:
    #     print(f"Error: {e}")
    # finally:
    #     cursor.close()
    #     conn.close()

    
